"""
@description: Entry point script to start the LLM fine-tuning process using configuration files.
@author: yangkw
@date: 2025-05-02
"""
# -*- coding: utf-8 -*-
import logging
import os
import sys
import json # 导入 json 用于加载 deepspeed config
from pathlib import Path
import argparse
from datetime import datetime # 导入 datetime
from accelerate import Accelerator # 导入 Accelerator
from accelerate.utils import DistributedDataParallelKwargs # 导入 DDP Kwargs

# 将 src 目录添加到 Python 路径，以便导入自定义模块
# 通常建议使用更规范的方式管理项目路径，例如安装为包或使用 PYTHONPATH
# 但对于简单脚本，这种方式比较直接
project_root = Path(__file__).parent.parent
src_path = project_root / "src"
sys.path.insert(0, str(project_root)) # 将项目根目录加入 sys.path

from transformers import HfArgumentParser, TrainingArguments as HfTrainingArguments # 导入官方 TrainingArguments

# 从 src 导入我们定义的模块和类
from src.utils import (
    logger, # 使用 utils 中配置好的 logger
    parse_args_from_yaml,
    ModelArguments,
    DataTrainingArguments,
    QuantizationArguments,
    PeftArguments,
    WandbArguments
)
from src.trainer import run_training

def main():
    """
    脚本主入口函数。
    解析命令行参数，加载配置，并启动训练流程。
    """
    # --- 1. 解析命令行参数 ---
    # 我们只接受一个必需参数：配置文件的路径
    parser = argparse.ArgumentParser(description="使用 YAML 配置文件启动大模型微调训练")
    parser.add_argument(
        "--config",
        type=str,
        required=True,
        help="指向 training_args.yaml 配置文件的路径。"
    )
    # HfArgumentParser 可以直接从 dataclass 创建参数，但我们这里选择从 YAML 加载
    # parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments, ...))
    # model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    cli_args = parser.parse_args()
    config_path = Path(cli_args.config)

    if not config_path.is_file():
        logger.error(f"指定的配置文件路径无效或文件不存在: {config_path}")
        sys.exit(1) # 配置文件不存在则退出

    logger.info(f"使用配置文件: {config_path}")

    # --- 2. 从 YAML 加载配置 ---
    try:
        # 使用我们 utils 中的函数解析 YAML 到定义的 dataclass
        # 注意：parse_args_from_yaml 返回的是我们自定义的 TrainingArguments dataclass
        (model_args, data_args, custom_training_args,
         quant_args, peft_args, wandb_args) = parse_args_from_yaml(config_path)
    except (FileNotFoundError, ValueError, TypeError) as e:
        logger.error(f"加载或解析配置文件 '{config_path}' 时出错: {e}")
        sys.exit(1)
    except Exception as e:
        logger.error(f"加载配置文件时发生未知错误: {e}")
        sys.exit(1)


    # --- 3. 初始化 Accelerator FIRST ---
    # 配置 DDP 参数，禁用未使用参数检查 (如果需要)
    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
    # 初始化 Accelerator，可以从环境或配置文件自动推断设置，或显式设置
    # 这里我们显式设置 bf16，并传入 ddp 配置
    accelerator = Accelerator(
        mixed_precision='bf16' if custom_training_args.bf16 else ('fp16' if custom_training_args.fp16 else 'no'), # 根据 YAML 配置设置精度
        kwargs_handlers=[ddp_kwargs]
    )
    logger.info(f"Accelerator initialized. State: {accelerator.state}")
    logger.info(f"Using {accelerator.num_processes} processes.")
    logger.info(f"Mixed precision: {accelerator.mixed_precision}")


    # --- 3.1. 确定 Output Directory (基于是否恢复训练) ---
    resume_checkpoint_path_str = custom_training_args.resume_from_checkpoint
    effective_output_dir_str = None
    final_model_dir_str = None # 最终模型保存目录

    if resume_checkpoint_path_str:
        logger.info(f"检测到 resume_from_checkpoint 设置: {resume_checkpoint_path_str}")
        resume_checkpoint_path = Path(resume_checkpoint_path_str)
        # 检查路径是否存在且看起来像一个检查点目录
        # 使用 Path.resolve() 获取绝对路径以进行更可靠的检查
        try:
            # 注意：resolve() 在路径不存在时会报错，我们需要先检查 is_dir()
            if resume_checkpoint_path.is_dir() and resume_checkpoint_path.name.startswith("checkpoint-"):
                 # 获取绝对路径以确保一致性
                abs_resume_path = resume_checkpoint_path.resolve()
                # 使用检查点所在的目录作为本次运行的输出目录
                effective_output_dir = abs_resume_path.parent
                effective_output_dir_str = str(effective_output_dir)
                logger.info(f"恢复训练，将使用原始输出目录: {effective_output_dir_str}")
                # 主进程确保目录存在 (理论上应该存在)
                if accelerator.is_main_process:
                     effective_output_dir.mkdir(parents=True, exist_ok=True)
            else:
                 logger.warning(f"指定的 resume_from_checkpoint 路径 '{resume_checkpoint_path_str}' 不是有效的检查点目录。")
                 resume_checkpoint_path_str = None # 清除路径以触发新目录逻辑
        except Exception as e: # 其他可能的错误，例如权限问题
             logger.error(f"检查 resume_from_checkpoint 路径时出错 '{resume_checkpoint_path_str}': {e}")
             resume_checkpoint_path_str = None # 清除路径以触发新目录逻辑

        # 如果路径无效，清除 custom_training_args 中的值，防止传递给 Trainer
        if not resume_checkpoint_path_str:
             logger.warning("将回退到创建新的时间戳目录。请检查 YAML 配置。")
             custom_training_args.resume_from_checkpoint = None

    # 如果不是恢复训练或恢复路径无效，则创建新的时间戳目录
    if not effective_output_dir_str:
        logger.info("未指定恢复检查点或路径无效，将创建新的时间戳输出目录。")
        timestamp_container = [None] # Use a list to hold the timestamp string
        if accelerator.is_main_process:
            timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
            timestamp_container[0] = timestamp
            logger.info(f"[Main Process] Generated timestamp: {timestamp}")
            # 主进程创建基础 output 目录 (如果不存在)
            (project_root / "output").mkdir(parents=True, exist_ok=True)

        # Broadcast the list containing the timestamp string from the main process to all processes
        import torch.distributed as dist
        dist.broadcast_object_list(timestamp_container, src=0) # src=0 means broadcast from rank 0

        # All processes now have the same timestamp
        timestamp = timestamp_container[0]
        if timestamp is None:
             logger.error("Failed to broadcast timestamp. Exiting.")
             sys.exit(1)

        if not accelerator.is_main_process:
            logger.info(f"[Process {accelerator.process_index}] Received timestamp: {timestamp}")

        # All processes construct the paths using the same timestamp
        model_path = Path(model_args.model_name_or_path)
        model_base_name = model_path.name
        run_output_dir = project_root / "output" / f"{model_base_name}_{timestamp}"
        effective_output_dir_str = str(run_output_dir) # 新的时间戳目录作为输出目录
        logger.info(f"[Process {accelerator.process_index}] 开始新训练，使用输出目录: {effective_output_dir_str}")
        # 主进程创建实际的运行输出目录
        if accelerator.is_main_process:
            run_output_dir.mkdir(parents=True, exist_ok=True)


    # 根据最终确定的输出目录计算 final_model_dir
    final_model_dir = Path(effective_output_dir_str) / "final_model"
    final_model_dir_str = str(final_model_dir)

    logger.info(f"[Process {accelerator.process_index}] 检查点和日志将保存在: {effective_output_dir_str}")
    logger.info(f"[Process {accelerator.process_index}] 最终模型将保存到: {final_model_dir_str}")


    # --- 3.2 确定 DeepSpeed 配置路径 (所有进程都需要知道路径) ---
    deepspeed_config_path_str = None
    # 优先使用 YAML 中指定的 deepspeed_config 路径
    if custom_training_args.deepspeed_config:
        ds_config_path_from_yaml = Path(custom_training_args.deepspeed_config)
        # 如果是相对路径，则相对于项目根目录或配置文件目录解析
        if not ds_config_path_from_yaml.is_absolute():
             # 尝试相对于项目根目录
             ds_config_path = project_root / ds_config_path_from_yaml
             if not ds_config_path.is_file():
                 # 尝试相对于配置文件目录
                 ds_config_path = config_path.parent / ds_config_path_from_yaml

        else:
            ds_config_path = ds_config_path_from_yaml

        if ds_config_path.is_file():
            deepspeed_config_path_str = str(ds_config_path)
            logger.info(f"使用 YAML 中指定的 DeepSpeed 配置文件: {deepspeed_config_path_str}")
        else:
            logger.warning(f"YAML 中指定的 DeepSpeed 配置文件未找到: {custom_training_args.deepspeed_config} (解析为 {ds_config_path})")
            # 可以选择在这里退出或继续而不使用 DeepSpeed
            # sys.exit(1)

    # 如果 YAML 中未指定，则尝试默认路径
    if not deepspeed_config_path_str:
        default_deepspeed_config_path = config_path.parent / "deepspeed_config.json"
        if default_deepspeed_config_path.is_file():
            deepspeed_config_path_str = str(default_deepspeed_config_path)
            logger.info(f"使用默认 DeepSpeed 配置文件: {deepspeed_config_path_str}")
        else:
             logger.info(f"未找到 DeepSpeed 配置文件 (YAML未指定，默认路径 {default_deepspeed_config_path} 也不存在)。将不使用 DeepSpeed。")


    # --- 3.3 处理报告目标 ---
    report_to = custom_training_args.report_to
    if report_to is None:
        report_to = "none" # 或者根据需要设置默认值，如 "all" 或 ["tensorboard"]
    elif isinstance(report_to, str): # 如果 YAML 里就是字符串
         pass # Trainer 能处理字符串
    elif isinstance(report_to, list):
         pass # Trainer 也能处理列表
    else:
        logger.warning(f"无法识别的 report_to 格式: {report_to}，将使用 'none'。")
        report_to = "none"


    # --- 3.4 创建官方 TrainingArguments 实例 ---
    # 处理早停逻辑
    load_best_at_end = custom_training_args.load_best_model_at_end
    metric_for_best = custom_training_args.metric_for_best_model
    greater_is_better_flag = custom_training_args.greater_is_better

    if custom_training_args.early_stopping_patience > 0:
        load_best_at_end = True # 启用早停时，强制加载最佳模型
        logger.info(f"早停已启用 (patience={custom_training_args.early_stopping_patience})，将自动设置 load_best_model_at_end=True。")
        # 如果 metric_for_best_model 未在 YAML 中指定，则默认为 eval_loss
        if metric_for_best is None:
             metric_for_best = "eval_loss"
             logger.info("metric_for_best_model 未指定，默认为 'eval_loss'。")
        # 如果 greater_is_better 未指定，根据 metric 推断 (loss 通常越小越好)
        if greater_is_better_flag is None:
             if "loss" in metric_for_best.lower():
                 greater_is_better_flag = False
                 logger.info(f"greater_is_better 未指定，根据指标 '{metric_for_best}' 推断为 False。")
             else:
                 greater_is_better_flag = True
                 logger.info(f"greater_is_better 未指定，根据指标 '{metric_for_best}' 推断为 True。")
    elif load_best_at_end and metric_for_best is None:
         # 如果只设置了 load_best_model_at_end=True 但没设置 metric，也需要默认值
         metric_for_best = "eval_loss"
         logger.info("load_best_model_at_end=True 但 metric_for_best_model 未指定，默认为 'eval_loss'。")
         if greater_is_better_flag is None:
             greater_is_better_flag = False
             logger.info(f"greater_is_better 未指定，根据指标 '{metric_for_best}' 推断为 False。")


    # 使用上面逻辑确定的 effective_output_dir_str 和从 YAML 读取的 resume_from_checkpoint
    # 以及处理后的早停相关参数
    training_args = HfTrainingArguments(
        output_dir=effective_output_dir_str, # 使用计算得到的输出目录
        resume_from_checkpoint=custom_training_args.resume_from_checkpoint, # 从 YAML 读取 (可能被置 None)
        num_train_epochs=custom_training_args.num_train_epochs,
        per_device_train_batch_size=custom_training_args.per_device_train_batch_size,
        per_device_eval_batch_size=custom_training_args.per_device_eval_batch_size,
        gradient_accumulation_steps=custom_training_args.gradient_accumulation_steps,
        learning_rate=float(custom_training_args.learning_rate), # 确保是 float
        lr_scheduler_type=custom_training_args.lr_scheduler_type,
        warmup_ratio=custom_training_args.warmup_ratio,
        weight_decay=custom_training_args.weight_decay,
        optim=custom_training_args.optim,
        logging_steps=custom_training_args.logging_steps,
        save_strategy=custom_training_args.save_strategy,
        save_steps=custom_training_args.save_steps,
        save_total_limit=custom_training_args.save_total_limit,
        eval_strategy=custom_training_args.evaluation_strategy,
        eval_steps=custom_training_args.eval_steps,
        seed=custom_training_args.seed,
        bf16=custom_training_args.bf16, # 从 YAML 读取
        fp16=custom_training_args.fp16, # 从 YAML 读取
        gradient_checkpointing=custom_training_args.gradient_checkpointing,
        report_to=report_to,
        deepspeed=deepspeed_config_path_str, # 传递 DeepSpeed 配置路径给 Trainer
        max_grad_norm=custom_training_args.max_grad_norm,
        # 早停相关参数
        load_best_model_at_end=load_best_at_end,
        metric_for_best_model=metric_for_best,
        greater_is_better=greater_is_better_flag,
        # 其他参数...
        remove_unused_columns=False, # 假设数据处理已完成列选择
        dataloader_num_workers = os.cpu_count() // accelerator.num_processes if os.cpu_count() else 4 # 调整 worker 数量
    )

    # --- 3.5 初始化早停回调 (如果启用) ---
    early_stopping_callback = None
    if custom_training_args.early_stopping_patience > 0:
        from transformers import EarlyStoppingCallback
        early_stopping_callback = EarlyStoppingCallback(
            early_stopping_patience=custom_training_args.early_stopping_patience
        )
        logger.info(f"已配置 EarlyStoppingCallback (patience={custom_training_args.early_stopping_patience})")

    # --- 4. 调用训练函数 ---
    trainer = None # 初始化 trainer 变量
    try:
        # 修改 run_training 调用，期望它返回 trainer 对象
        # 并将早停回调传递给它 (如果已创建)
        trainer = run_training(
            model_args=model_args,
            data_args=data_args,
            training_args=training_args, # 传递官方的 TrainingArguments
            quant_args=quant_args,
            peft_args=peft_args,
            wandb_args=wandb_args,
            overwrite_cache=custom_training_args.overwrite_cache, # 传递 overwrite_cache 值
            early_stopping_callback=early_stopping_callback # 传递早停回调
        )
        logger.info("训练流程 run_training 执行完成。")

        # --- 5. 保存最终模型到指定目录 ---
        if trainer is not None:
            if accelerator.is_main_process:
                logger.info(f"[主进程] 准备保存最终模型到: {final_model_dir_str}")
                # 主进程负责创建最终模型目录 (如果尚不存在)
                Path(final_model_dir_str).mkdir(parents=True, exist_ok=True)

            # 所有进程都需要调用 save_model，因为这可能是一个涉及集合通信的操作 (例如 ZeRO Stage 3)
            # Trainer 和 Accelerate 会在内部处理好只让主进程实际写入文件
            logger.info(f"[进程 {accelerator.process_index}] 调用 trainer.save_model() 保存到 {final_model_dir_str}")
            trainer.save_model(final_model_dir_str)

            # 同样，保存 tokenizer 也应该由所有进程调用 (如果它也涉及分布式状态)
            # 或者至少在 Trainer 内部正确处理
            if hasattr(trainer, 'tokenizer') and trainer.tokenizer is not None:
                logger.info(f"[进程 {accelerator.process_index}] 调用 trainer.tokenizer.save_pretrained() 保存到 {final_model_dir_str}")
                trainer.tokenizer.save_pretrained(final_model_dir_str)
            
            if accelerator.is_main_process:
                logger.info(f"[主进程] 模型和 Tokenizer (如果适用) 的保存调用已在所有进程上发起。")
        else:
            # 只在主进程记录 Trainer 为 None 的警告
            if accelerator.is_main_process:
                 logger.warning("[主进程] Trainer 对象为 None 或训练失败，无法保存最终模型。")

        # 等待所有进程完成所有操作，包括可能的分布式保存
        logger.info(f"[进程 {accelerator.process_index}] 等待所有进程同步...")
        accelerator.wait_for_everyone()

        if accelerator.is_main_process:
            logger.info("所有进程完成，训练脚本执行成功完成。")

    except Exception as e:
        logger.exception(f"训练过程中发生未捕获的错误: {e}") # 使用 exception 记录堆栈信息
        sys.exit(1) # 发生错误时退出

if __name__ == "__main__":
    main()

# 如何运行:
# 您需要使用 `accelerate launch` 来启动这个脚本，以便正确配置分布式环境（即使是单 GPU 也建议使用）。
#
# 例如，假设您有 2 个 GPU：
# accelerate launch \
#     --config_file accelerate_config.yaml \ # accelerate 的配置文件 (如果使用)
#     --num_processes 2 \                     # 使用的 GPU 数量
#     scripts/run_train.py \
#     --config config/training_args.yaml      # 我们项目的训练配置文件
#
# 如果只有 1 个 GPU：
# accelerate launch \
#     --num_processes 1 \
#     scripts/run_train.py \
#     --config config/training_args.yaml
#
# `accelerate` 会自动读取 `training_args.yaml` 中指定的 `deepspeed` 配置（通过 HfTrainingArguments 的 `deepspeed` 参数传递），
# 并设置相应的环境变量和启动参数。
